
import string
from binascii import hexlify
'''
As the generators analyze an input stream they will create instances of these primitives
These classes will then be feed to the fuzz script writer
The script writer will replace the parts of the stream the generators have found to match primitives
'''

class primitiveBase(object):
    '''
    classdocs
    '''
    typeSize = -1 # 0 for block definitions and gt 0 for everything else
    originalValue = -1
    streamOffset = -1
    name = ""

    def getName(self):
        return self.name
    
    #given the name of a lego, returns the name of the block that surrounds it
    def getBlockName(self):
        return "BLOCK_%s" % (self.name)

    def getStart(self):
        return self.streamOffset
    
    def getEnd(self):
        return self.streamOffset + self.typeSize
    
    def getSulleyString(self):
        return "ERROR"
    
    def __init__(self):
        pass
    
    def __str__(self):
        return self.getSulleyString()
    
    def __cmp__(self, other):
        if self.streamOffset == other.streamOffset:
            if self.typeSize == 0:
                if other.typeSize == 0: #two size 0 primitives at the same offset (two blocks start/end at the same offset)
                    return 0
                else: #a size 0 primitive comes before others (think block definitiions)
                    return -1
            else:
                if other.typeSize == 0:
                    return 1
                else:
                    raise Exception, "Bug detected, overlapping primitives found"
                    
        elif self.streamOffset > other.streamOffset:
            if other.isStreamRegionOwner(self.streamOffset):
                print "raising exception"
                raise Exception, "Bug detected, overlapping primitives found" #if we are sorting a list of primitives and this happens we constructed the list incorrectly
            else:
                return 1
        else:
            if self.isStreamRegionOwner(other.streamOffset):
                print "raising exception"
                raise Exception, "Bug detected, overlapping primitives found"
            else:
                return -1
        raise Exception, "WTF?"
                
    def __lt__(self, other): 
        return self.__cmp__(other) < 0
    
    def isStreamRegionOwner(self, streamOffset, length=1):
        if length < 1:
            raise Exception, "Cannot determine owner of 0 length region"
        #if typeSize is 0, always return false
        if streamOffset >= self.streamOffset and (streamOffset + length) <= (self.streamOffset + self.typeSize):
            return True
        return False
    
    
    #escape metachars that would cause problems in the resulting python script
    #leave basic chars as-is so the script is more readable.
    @staticmethod
    def escapeForPython(data):
        encoded = []
        for c in data:
            if c in string.ascii_letters or c in string.digits:
                encoded.append(c)
            else:
                encoded.append("\\x%s" % (hexlify(c)))
        s = "".join(encoded)
        return s
        
        

class static(primitiveBase):
    
    def __init__(self, value="", streamOffset=-1):
        self.originalValue = value
        self.streamOffset = streamOffset
        self.typeSize = len(value)

    def getSulleyString(self):
        s = """s_static("%s")\n""" % self.escapeForPython(self.originalValue)
        return s
        
    #given a set of locations, creates statics for all other locations.
    #returns a list of statics
    @staticmethod
    def buildAllStatics(fuzzedLocations, inputBuff):
        tempSet = set(range(0, len(inputBuff)))
        staticLocations = tempSet - fuzzedLocations
        staticLocations = list(staticLocations)
        staticLocations.sort() #ensure order is as expected
        start = -1
        end = -1
        statics = []
        for x in staticLocations:
            if start == -1:
                start = x
            if x+1 in staticLocations:
                continue
            else:
                end = x
                if end < start:
                    raise Exception, "Bug detected in buildAllStatics"
                newStatic = static(inputBuff[start:end+1], start)
                statics.append(newStatic)
                start = -1
        return statics
            
            
class block(primitiveBase):
    typeSize = 0
    blockStart = None    
    SULLEYFUNCTION = "s_block"
    def __init__(self, blockName="", streamOffset=-1, blockStart=-1):
        self.streamOffset = streamOffset
        self.blockStart = blockStart
        self.name = blockName
    
    def getSulleyString(self):
        if (self.blockStart):
            tempStr = "_start"
        else:
            tempStr = "_end"
        s = """%s%s("%s")\n""" % (self.SULLEYFUNCTION, tempStr, self.name)
        return s
        
    def createBlocks(self, stringLego):
        if type(stringLego) == stringPrimitive:
            startBlock = block(stringLego.getBlockName(), stringLego.getStart(), True)
            endBlock = block(stringLego.getBlockName(), stringLego.getEnd(), False)
            return (startBlock, endBlock)
        else:
            return ()
        
class word16(primitiveBase):
    endian = "<"
    name = ""
    #blockStart = -1
    #blockEnd = -1
    unpackFormats = ("h", "H") #note could remove the upper case unpack format since its unlikely we will ever search for negative numbers
    typeSize = 2 #16bits
    format = "binary"
    signed = False
    SULLEYFUNCTION = "s_word"
    variance = 0
    
    def __init__(self, originalValue=-1, streamOffset=-1, endian="<", name="", format="binary", signed=False, variance=0):
        self.endian = endian
        self.name = name
        self.streamOffset = streamOffset
        self.originalValue = originalValue
        self.format = format
        self.signed = signed
        self.variance = variance
        
    def getSulleyString(self):
        s = """%s(%d, endian="%s", format="%s", signed=%s, name="%s")\n""" % (self.SULLEYFUNCTION, self.originalValue, self.endian, self.format, self.signed, self.name)
        return s  
      
class word32(word16):
    unpackFormats = ("i", "I")
    typeSize = 4 #32bits
    SULLEYFUNCTION = "s_int"
        
class word64(word16):
    unpackFormats = ("q", "Q")
    typeSize = 8 #64bits
    SULLEYFUNCTION = "s_double"

class sizer(word32, word64, word16):
    SULLEYFUNCTION = "s_size"
    blockName = None
    variance = 0

    #does the sizer need to be aware of the location of the block?
    #sizer can be based off an already existent word
    def __init__(self, wordObject):
        #self = wordObject ?? doesnt work ??
        self.endian = wordObject.endian
        self.name = wordObject.name
        self.streamOffset = wordObject.streamOffset
        self.originalValue = wordObject.originalValue
        self.format = wordObject.format
        self.signed = wordObject.signed
        self.typeSize = wordObject.typeSize
        self.variance = wordObject.variance
   
    def setBlockName(self, blockName): 
        self.blockName = blockName
        
    def getSulleyString(self):
        if self.blockName == None:
            raise Exception, "blockName must be set for sizer"
        lambdaDef = """lambda_%s = lambda x: x + %d\n""" % (self.blockName, self.variance)
        s = """%s("%s", length=%d, endian="%s", format="%s", signed=%s, name="%s", fuzzable=True, math=lambda_%s)\n""" % (self.SULLEYFUNCTION, self.blockName,  self.typeSize, self.endian, self.format, self.signed, self.name, self.blockName)
        return lambdaDef + s

class stringPrimitive(primitiveBase):
    numericSizer = False
    SULLEYFUNCTION = "s_string"
    encoding = "ascii"
    
    def __init__(self, originalValue, streamOffset, encoding="ascii", numericSizer = False, name=""):
        self.numericSizer = numericSizer
        self.typeSize = len(originalValue)
        self.streamOffset = streamOffset
        self.originalValue = originalValue
        self.encoding = encoding
        self.name = name

    def getSulleyString(self):
        s = """%s("%s", encoding="%s")\n""" % (self.SULLEYFUNCTION, self.escapeForPython(self.originalValue), self.encoding)
        return s
    
      
